"""
https://www.tensorflow.org/guide/autodiff
"""

import sys
from python_ai.common.xcommon import *
import tensorflow as tf

print(tf.__version__)
tf.random.set_seed(777)

sep('y = x**2')
x = tf.Variable(3.0, name='x')
with tf.GradientTape() as tape:
    y = x**2
dy_dx = tape.gradient(y, x)
sep('dy_dx')
print(dy_dx)
sep('watched vars')
print([(v.name, v.value()) for v in tape.watched_variables()])

sep('y = x@w + b')
x = [[1., 2., 3.]]
w = tf.Variable(tf.random.normal([3, 2]), name='w')
b = tf.Variable(tf.random.normal([1, 2]), name='b')
with tf.GradientTape() as tape:
    y = x @ w + b
    loss = tf.reduce_mean(y ** 2)
sep('by list')
[dl_dw, dl_db] = tape.gradient(loss, [w, b])
print(dl_dw)
print(dl_db)
sep('by dict')
with tf.GradientTape() as tape:
    y = x @ w + b
    loss = tf.reduce_mean(y ** 2)
dict = tape.gradient(loss, {'w': w, 'b': b})
print(dict)
sep('watched vars')
print([(v.name, v.value()) for v in tape.watched_variables()])
